#include "ParticleSystem.h"
#include "ParticleSystem.cuh"

ParticleSystem::ParticleSystem(int size)
{
    CglShader shaderVS, shaderPS, shaderShadowPS, shaderPassThruVS, shaderBlurPS, shaderFloorVS, shaderFloorPS, shaderTex2dPS;

    m_size = size;
    m_numSlices = 32;

    m_numSlices = 32;
    m_sliceNo = 0;

    m_pointRadius = 1;
    m_pointScale  = 1;
    m_shadowAlpha = 0.02;
    m_spriteAlpha = 0.1;
    m_emitIndex   = 0;

    m_Pos.Create   (size, BUFFER_H_AND_B);
    m_Vel.Create   (size, BUFFER_H_AND_B);
    m_Color.Create (size, BUFFER_H_AND_B);
    m_Zs.Create    (size, BUFFER_H_AND_B);

    m_sliceNo = m_numSlices - 1;

    set(m_Position, 0.0f, 10.0f, 0.0f, 1.0f);
    set(m_Velocity, 0.0f,  0.0f, 0.0f, 0.0f);

    // load shaders
    shaderVS.CreateFromFile("..\\Media\\shaders\\particle.vs", GL_VERTEX_SHADER);
    shaderPS.CreateFromFile("..\\Media\\shaders\\particle.ps", GL_FRAGMENT_SHADER);
    shaderShadowPS.CreateFromFile("..\\Media\\shaders\\particleShadow.ps", GL_FRAGMENT_SHADER);

    m_particleProg.Create  ( shaderVS.Id(), shaderPS.Id()       );
    m_shadowProg.Create    ( shaderVS.Id(), shaderShadowPS.Id() );

    m_pShadow = NULL;
    m_pHalfRes = NULL;
}

ParticleSystem::~ParticleSystem()
{
    m_particleProg.Delete();
    m_shadowProg.Delete();

    m_Pos.Release(BUFFER_H_AND_B);
    m_Vel.Release(BUFFER_H_AND_B);
    m_Color.Release(BUFFER_H_AND_B);
    m_Zs.Release(BUFFER_H_AND_B);

    m_pShadow = NULL;
    m_pHalfRes = NULL;

    m_plCamera = NULL;
    m_pvCamera = NULL;
}

void ParticleSystem::Sync(C_Camera *pCamera)
{
    set(m_Position, pCamera->Eye(), 1.0f);
    set(m_Velocity, pCamera->Dir(), 0.1f);
}

void ParticleSystem::Reset()
{
    m_numActive = 0;
    m_emitIndex = 0;
}

void ParticleSystem::SetOffscreenBufferPtr(glLightMap * pShadow, glLightMap * pHalfRes)
{
    m_pShadow  = pShadow;
    m_pHalfRes = pHalfRes;
}

void ParticleSystem::InitSphere(int3 dim, float r, bool bCUDA)
{
    int i = 0;
    for(int z=0; z < dim.z; z++) 
    {
        for(int y=0; y<dim.y; y++) 
        {
            for(int x=0; x<dim.x; x++) 
            {
                float4 p;
                p.x = (float)x / dim.x;
                p.y = (float)y / dim.y;
                p.z = (float)z / dim.z;
                p.w = 0.0f;

                p = (p * 2.0f - 1.0f) * r;
 
                float d = length3(p);

                if (i < m_size)
                {
                    if ((d < r)) 
                    {
                        m_Pos[0][i] = this->m_Position + p;
                        m_Vel[0][i] = this->m_Velocity;
                        set(m_Color[0][i], rnsFloat3(), 1.0f);
                    }
                    else
                    {
                        float3 randomPos = xyz(this->m_Position) + rnSphere() * r;

                        set(m_Pos[0][i], randomPos, 1.0f);
                        m_Vel[0][i] = this->m_Velocity;
                        set(m_Color[0][i], rnsFloat3(), 1.0f);
                    }
                }

                i++;
            }
        }
    }
    m_numActive = i;
}

void ParticleSystem::InitGrid(int3 dim, float r, bool bCUDA)
{
    int i = 0;
    for(int z=0; z<dim.z; z++) 
    {
        for(int y=0; y<dim.y; y++) 
        {
            for(int x=0; x<dim.x; x++) 
            {
                float4 p;
                p.x = (float)x / dim.x;
                p.y = (float)y / dim.y;
                p.z = (float)z / dim.z;
                p.w = 0.0f;

                p = (p * 2.0f - 1.0f) * r;

                if (i < m_size) 
                {
                    m_Pos[0][i] = this->m_Position + p;
                    m_Vel[0][i] = this->m_Velocity;
                    set(m_Color[0][i], rnsFloat3(), 1.0f);
                    i++;
                }
            }
        }
    }
    m_numActive = i;
}

// add some procedural noise
void ParticleSystem::AddNoisePosition(float freq, float scale, bool bCUDA)
{
    if (bCUDA)
    {
        m_Pos.Map(BUFFER_R);

        m_Pos.Memcpy(BUFFER_H_TO_D, BUFFER_R);

        ParticleSystem_AddPositionNoise( m_Pos(BUFFER_R), freq, scale, m_Pos.size() );

        m_Pos.Memcpy(BUFFER_D_TO_H, BUFFER_R);

        m_Pos.Unmap(BUFFER_R);
    }
    else
    {
        for(int i=0; i<m_size; i++) 
        {
            float3 pn = ImprovedPerlinNoise::BrownianMotionf3( xyz(m_Pos[0][i]) * freq ) * scale;

            float3 new_pos = xyz(m_Pos[0][i]) + pn;
            set(m_Pos[0][i], new_pos, 1.0);
        }
    }
}

void ParticleSystem::AddNoiseVelocity(float freq, float scale, float3 offset, int sett, bool bCUDA)
{
    for(int i=0; i<m_size; i++) 
    {
        //if ((i & 7) == sett) 
        //{
        float3 pn = ImprovedPerlinNoise::BrownianMotionf3(xyz(m_Pos[0][i])*freq+offset)*scale;
        set(m_Vel[0][i], pn, 0.0f);
        //}
    }
}

void ParticleSystem::Animate(bool bCUDA)
{
    const float4 gravity = {0.0f, -0.001f, 0.0f, 0.0f};
    const float  damping = 0.98;

    if (bCUDA)
    {
        m_Pos.Map(BUFFER_R);
        m_Vel.Map(BUFFER_R);

        m_Pos.Memcpy(BUFFER_H_TO_D, BUFFER_R);
        m_Vel.Memcpy(BUFFER_H_TO_D, BUFFER_R);

        ParticleSystem_Animate( m_Pos(BUFFER_R), m_Vel(BUFFER_R), gravity, damping, m_numActive );

        m_Pos.Memcpy(BUFFER_D_TO_H, BUFFER_R);
        m_Vel.Memcpy(BUFFER_D_TO_H, BUFFER_R);

        m_Pos.Unmap(BUFFER_R);
        m_Vel.Unmap(BUFFER_R);
    }
    else
    {
        // integrate
        for(int i=0; i<m_numActive; i++) 
        {
            m_Pos[0][i] += m_Vel[0][i];
            m_Vel[0][i] *= damping;
            m_Vel[0][i] += gravity;

            // TODO : set proper bnd condition 

            if (m_Pos[0][i].y < -1.0) 
            {
                m_Pos[0][i].y = -1.0;
                m_Vel[0][i].y *= -0.9;
            }
        }
    }
}

void ParticleSystem::Emit(float speed, bool bCUDA)
{
    const int emitRate = 50;
    const int clumpSize = 10;
    float angle = 0.1f;

    // emitter
    for(int i=0; i<emitRate; i++) {

        float3 basePos = xyz(this->m_Position) + rnSphere() * f3(0.2, 0.0, 0.2);
        float3 dir = normalize( rnSphere() );
        float3 baseVel = normalize(xyz(this->m_Velocity) + dir * angle) * speed;

        for(int j=0; j<clumpSize; j++)
        {
            set( m_Pos[0][m_emitIndex], basePos + rnSphere() * 0.05f  );
            set( m_Vel[0][m_emitIndex], baseVel + rnsFloat3() * 0.002f );

            m_emitIndex = (m_emitIndex + 1) % m_size;
        };

        if (m_emitIndex > m_numActive) 
        {
            m_numActive = m_emitIndex;
        }
    }
}

void ParticleSystem::glRenderPoints(int start, int count)
{
    if (count == 0)
        count = m_size;

    glPointSize(5);
    glBegin(GL_POINTS);
    for(int i=start; i<start + count; i++) 
    {
        glVertex3fv((float *) &(m_Pos[0][i]));
    }
    glEnd();
    glPointSize(1);
}

void ParticleSystem::glRenderSorted(int start, int count, float alpha)
{
    uint *sortedIndices = sorter.GetIndices();

    glBegin(GL_POINTS);
    for(int i=start; i<start + count; i++) 
    {
        int si = sortedIndices[i];

        glTexCoord3fv((float *) &m_Vel[0][si]);
        glVertex3fv  ((float *) &m_Pos[0][si]);
    }
    glEnd();
}

void ParticleSystem::glRenderSprites(CglProgram * pProg, 
                                     int start, int count, 
                                     float alpha, float fov)
{
    glEnable(GL_POINT_SPRITE_ARB);
    glTexEnvi(GL_POINT_SPRITE_ARB, GL_COORD_REPLACE_ARB, GL_TRUE);
    glEnable(GL_VERTEX_PROGRAM_POINT_SIZE_NV);

    glEnable(GL_DEPTH_TEST);
    glDepthMask(GL_FALSE);  // don't write depth

    glEnable(GL_BLEND);

    GLint viewport[4];
    glGetIntegerv(GL_VIEWPORT, viewport);

    float winPointScale = viewport[3] / tan(to_radian(fov * 0.5f));
    float invW = 1.0f / viewport[2];
    float invH = 1.0f / viewport[3];

    pProg->Enable();

    pProg->setUniform1f( "pointRadius",   m_pointRadius                 );
    pProg->setUniform1f( "pointScale",    m_pointRadius * winPointScale );
    pProg->setUniformfv( "halfVectorEye", (float *) &m_halfDirES, 3     ); 
    pProg->setUniform2f( "invViewport",   invW, invH );

    pProg->bindTexture ( "shadowTex", m_pShadow->getColor(),  GL_TEXTURE_2D, 0 );
    pProg->bindTexture ( "depthTex",  m_pHalfRes->getDepth(), GL_TEXTURE_2D, 1 );

    glRenderSorted(start, count, alpha);

    pProg->Disable();

    glDisable(GL_POINT_SPRITE_ARB);
    glDisable(GL_BLEND);
    glDepthMask(GL_TRUE);
}


void ParticleSystem::glRenderView(int i)
{
    glMatrixMode(GL_MODELVIEW);
    glPushMatrix();
    glLoadMatrixf((GLfloat *) &m_vvMatrix);

    glMatrixMode(GL_PROJECTION);
    glPushMatrix();
    glLoadMatrixf((GLfloat *) &m_vpMatrix);

    m_pHalfRes->glBegin();

    glColor4f(1.0, 1.0, 1.0, m_spriteAlpha);

    if (m_invView) glBlendFunc(GL_ONE_MINUS_DST_ALPHA, GL_ONE); // front-to-back
    else           glBlendFunc(GL_ONE, GL_ONE_MINUS_SRC_ALPHA); // back-to-front
    
    glRenderSprites( &m_shadowProg, i*m_batchSize, m_batchSize, m_spriteAlpha, m_pvCamera->FoV());         

    m_pHalfRes->glEnd();

    glMatrixMode(GL_PROJECTION);
    glPopMatrix();
    glMatrixMode(GL_MODELVIEW);
    glPopMatrix();
}

// draw slice from light's point of view
void ParticleSystem::glRenderLight(int i)
{
    glMatrixMode(GL_MODELVIEW);
    glPushMatrix();
    glLoadMatrixf((GLfloat *) &m_lvMatrix);

    glMatrixMode(GL_PROJECTION);
    glPushMatrix();
    glLoadMatrixf((GLfloat *) &m_lpMatrix);

    m_pShadow->glBegin();    

    glColor4f(1.0, 1.0, 1.0, m_shadowAlpha);

    glBlendFunc(GL_SRC_ALPHA, GL_ONE_MINUS_SRC_ALPHA);

    glRenderSprites( &m_particleProg, i*m_batchSize, m_batchSize, m_shadowAlpha, m_plCamera->FoV());

    m_pShadow->glEnd();

    glMatrixMode(GL_PROJECTION);
    glPopMatrix();
    glMatrixMode(GL_MODELVIEW);
    glPopMatrix();
}

void ParticleSystem::glRenderBatchs()
{
    m_pShadow->glBegin(); 
    glClearColor(0.0, 0.0, 0.0, 0.0); 
    glClear(GL_COLOR_BUFFER_BIT);
    m_pShadow->glEnd();

    // clear volume image
    m_pHalfRes->glBegin();
    glClearColor(0.0, 0.0, 0.0, 0.0); 
    glClear(GL_COLOR_BUFFER_BIT);
    m_pHalfRes->glEnd();         

    // render slices
    for(int i=0; i<m_sliceNo+1; i++) 
    {
        // draw slice from camera view, sampling light buffer
        glRenderView(i);

        // draw slice from light view to light buffer, accumulating shadows
        glRenderLight(i);
    }
}

void displayTexture(GLuint tex, CglProgram * pDisplay)
{
    pDisplay->Enable();
    pDisplay->bindTexture("tex", tex, GL_TEXTURE_2D, 0);
    glRenderQuad();
    pDisplay->Disable();
}


// composite final volume image on top of scene
void ParticleSystem::Composite(uint32 w, uint32 h, CglProgram * pDisplay)
{
    glViewport(0, 0, w, h);
    glDisable(GL_DEPTH_TEST);
    glBlendFunc(GL_ONE, GL_ONE_MINUS_SRC_ALPHA);
    glEnable(GL_BLEND);

    displayTexture(m_pHalfRes->getColor(), pDisplay);

    glDisable(GL_BLEND);
}

void ParticleSystem::ComputeHalf()
{
    if (dot(m_viewDir, m_lightDir) > 0) 
    {
        set( m_halfDir, normalize( m_viewDir + m_lightDir), 0.0f );
        m_invView = false;
    } 
    else 
    {
        set( m_halfDir, normalize(-m_viewDir + m_lightDir), 0.0f );
        m_invView = true;
    }

    matrix4x4 modelView;
    glGetFloatv(GL_MODELVIEW_MATRIX, (float *) &modelView);

    Matrix4x4LookAtRHR( &m_lvMatrix, 
                         m_plCamera->Eye(), 
                         m_plCamera->At(), 
                         m_plCamera->Up() );

    Matrix4x4LookAtRHR( &m_vvMatrix, 
                         m_pvCamera->Eye(), 
                         m_pvCamera->At(), 
                         m_pvCamera->Up() );


    Matrix4x4PerspectiveFovGLR( &m_lpMatrix, 
                                 m_plCamera->FoV(), 
                                 m_plCamera->Asp(), 
                                 m_plCamera->zN(), 
                                 m_plCamera->zF());

    Matrix4x4PerspectiveFovGLR( &m_vpMatrix, 
                                 m_pvCamera->FoV(), 
                                 m_pvCamera->Asp(), 
                                 m_pvCamera->zN(), 
                                 m_pvCamera->zF());
    
    // construct shadow matrix
    matrix4x4 scale, translate, m_invMV;

    Matrix4x4Scale      ( &scale, 0.5, 0.5, 0.5, 1.0f);
    Matrix4x4TranslateR ( &translate, f3(0.5, 0.5, 0.5));
    Matrix4x4Inverse    ( &m_invMV, &m_vvMatrix );

    m_sMatrix = m_invMV * m_lvMatrix * m_lpMatrix * scale * translate;

    glActiveTexture(GL_TEXTURE0);
    glMatrixMode(GL_TEXTURE);
    glLoadMatrixf((GLfloat *) &m_sMatrix);

    // calc half vector in eye space
    m_halfDirES = m_vvMatrix * m_halfDir;

    SortByHalfDir();
}

void ParticleSystem::SetCameraPtr(C_Camera * pvCamera, C_Camera * plCamera)
{
    m_pvCamera = pvCamera;
    m_plCamera = plCamera;

    m_viewDir  =  pvCamera->Dir();
    m_lightDir = -plCamera->Dir(); // Camera and light dir are towards the scene
}

void ParticleSystem::glRender(C_Camera * pvCamera, C_Camera * plCamera)
{
    SetCameraPtr(pvCamera, plCamera);

    glRender();
}

void ParticleSystem::glRender()
{
    ComputeHalf();

    SortByHalfDir();

    glRenderBatchs(); //error
}

void ParticleSystem::SortByHalfDir()
{
    for(int i=0; i<m_numActive; i++) 
    {
        float z = -dot(m_halfDir, m_Pos[0][i]);  // project onto vector
        m_Zs[0][i] = z;
    }

    sorter.Sort(m_Zs[0], m_numActive);

    m_batchSize = m_numActive / m_numSlices;
}